﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;

namespace Algoritmia.Problems.Puzzles.NQueens
{
   
    public class NQueensStateSpace3 : IReversibleForwardStateSpace<NQueensStateSpace3.State,int>
    {
        private readonly int _n;

        public class State
        {
            public readonly State Prev;
            public readonly int Row;
            public readonly int N;
            private readonly bool[] _freeRow;
            private readonly bool[] _diag;
            private readonly bool[] _gaid;

            public State(int n,State prev = null, int row = 0)
            {
                if ( prev == null)
                {
                    N = 0;
                    _freeRow = new bool[n];
                    FillTrue(_freeRow);
                    _diag = new bool[2*n - 1];
                    FillTrue(_diag);
                    _gaid = new bool[2*n - 1];
                    FillTrue(_gaid);
                }
                else
                {
                    N = prev.N + 1;
                    _freeRow = prev._freeRow;
                    _diag = prev._diag;
                    _gaid = prev._gaid;
                    var col = N - 1;
                    
                    _freeRow[row] = _diag[col + row] = _gaid[ GaidIndex(col,row) ] = false;
                }
                Prev = prev;
                Row = row;
            }

            private static void FillTrue(bool[] array)
            {
                for(int i = 0; i < array.Length; i++)
                {
                    array[i] = true;
                }
            }

            public void RemoveLast()
            {
                var col = N - 1;
                var row = Row;

                _freeRow[row] = _diag[col + row] = _gaid[GaidIndex(col,row)] = true;
            }

            private int GaidIndex(int col, int row)
            {
                var gaidIndex = (col - row + N - 1);
                gaidIndex = gaidIndex < 0
                                    ? _gaid.Length + gaidIndex
                                    : gaidIndex;
                return gaidIndex;
            }

            public bool CellIsFree(int col, int row)
            {
                return _freeRow[row] && _diag[col+row] && _gaid[ GaidIndex(col,row)];
            }


            public Lists.IList<int> ToList()
            {
                State s = this;
                int[] array = new int[N];
                while (s.Prev != null)
                {
                    array[s.N - 1] = s.Row;
                    s = s.Prev;
                }
                return AlgoritmiaFactory.Lists<int>.BuildDefault(array);
            }
        }

        public NQueensStateSpace3(int N)
        {
            _n = N;
        }

        public State Decide(State state, int decision)
        {
            return new State(_n,state, decision);
        }

        public IEnumerable<int> Decisions(State state)
        {
            var col = state.N;
            for(int row = 0; row < _n; row++)
            {
                if ( state.CellIsFree(col,row) )
                {
                    yield return row;
                }
            }
        }

        public IEnumerable<State> InitialStates
        {
            get { yield return new State(_n); }
        }

        public bool IsFinal(State state)
        {
            return _n == state.N;
        }

        public State Undo(State state, int decision)
        {
            state.RemoveLast();
            return state.Prev;
        }
    }


    public class NQueensEnumerator3
    {
        protected readonly Func<Lists.IList<int>> _createList;

        public NQueensEnumerator3(Func<Lists.IList<int>> createList = null)
        {
            _createList = createList ?? (() => AlgoritmiaFactory.Lists<int>.Default);
        }

        protected virtual IEnumerable<NQueensStateSpace3.State> Backtracking(NQueensStateSpace3.State state,
            NQueensStateSpace3 space, Lists.IList<int> decisions)
        {
            if (space.IsFinal(state))
            {
                yield return state;
            }

            foreach (var decision in space.Decisions(state))
            {
                decisions.Add(decision);
                state = space.Decide(state, decision);
                foreach (var result in Backtracking(state, space, decisions))
                {
                    yield return result;
                }
                decisions.RemoveAt(decisions.Count - 1);
                state = space.Undo(state, decision);
            }
        }

        public IEnumerable<NQueensStateSpace3.State> Enumerate(NQueensStateSpace3 space)
        {
            var decisions = _createList();
            foreach(var initial in space.InitialStates)
            {
                foreach (var result in Backtracking(initial, space, decisions))
                {
                    yield return result;
                }
            }
        }

        public NQueensStateSpace3.State FirstSolution(NQueensStateSpace3 space)
        {
            return Enumerate(space).First();
        }
    }

}
